# 导入相关模块
import torch
from torch.utils.data import TensorDataset  # 构造数据集对象
from torch.utils.data import DataLoader  # 数据加载器
from torch import nn  # nn模块中有平方损失函数和假设函数
from torch import optim  # optim模块中有优化器函数
from sklearn.datasets import make_regression  # 创建线性回归模型数据集
import matplotlib.pyplot as plt  # 可视化

plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


##########################
# 使用pytorch搭建线性回归模型

#1.准备数据集
#2.构建模型
#3.设置损失函数和优化器（结合反向传播更新参数w.data = w.data - n*梯度）
#4.训练模型
#
#nn.MSELoss() 均方误差损失函数
#data.DataLoader() 数据加载器,模拟真实场景中批量获取数据
#optim.SGD() 随机梯度下降优化器
#nn.Linear() 线性回归模型   y=wx+b
##########################


# numpy对象 -> torch.tensor对象 -> 数据集对象TensorDataset -> 数据加载器对象DataLoader
#  -> 线性回归模型 -> 损失函数对象 -> 优化器对象 -> 训练模型


def create_dataset():
    # 创建数据集对象
    x , y , coef = make_regression(
        n_samples=100,  # 样本数100 相当于y有100个值
        n_features=1,  # 特征数 相当于x是一维的
        noise=20,  # 噪声   越大，数据越分散
        coef=True,  # 返回真实系数w
        bias=14.5,  # 偏置项b
        random_state=5 # 随机种子
    )
    print("真实系数w:", coef)
    print("x的形状:", x.shape)
    print("y的形状:", y.shape)

    #目前还是numpy对象
    x = torch.tensor(x, dtype=torch.float32)
    y = torch.tensor(y, dtype=torch.float32)

    return x , y ,coef

def train(x,y,coef):
    # 构造数据集对象 和 数据加载器对象
    dataset = TensorDataset(x,y)
    # shuffle=True 打乱数据集
    dataLoader = DataLoader(dataset,batch_size=16,shuffle=True)

    #创建模型
    model = nn.Linear(in_features=1,out_features=1)  # 线性回归模型,一个输入一个输出
    criterion = nn.MSELoss()  
    optimizer = optim.SGD(model.parameters(),lr=0.01) 

    epochs=100
    loss_list = [] 
    total_loss = 0.0 
    total_samples = 0

    for epoch in range(epochs):
        for train_x,train_y in dataLoader:
            Y_pred = model(train_x)
            loss = criterion(Y_pred,train_y.reshape(-1,1))

            total_loss += loss.item()
            total_samples += 1

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        loss_list.append(loss.item() / total_samples)
        print(f"第{epoch + 1}轮 ,平均损失:{loss_list[-1]:.4f}")    
    
    print("训练完成! ")
    print(f"{epochs}轮后，模型最终损失:{loss_list[-1]}")
    print(f"模型参数w:{model.weight.data.numpy()[0][0]:.3f},b:{model.bias.data.numpy()[0]:.3f}")

    # ============================================
    # 🔥 下面开始：画图（不改变你原来的训练代码）
    # ============================================

    # 原始散点图
    plt.scatter(x.numpy(), y.numpy(), color="blue", label="原始数据点")

    # 真实直线（coef + bias）
    true_bias = 14.5   # 你的 make_regression 生成时写死了
    x_line = x.numpy()
    y_true_line = coef * x_line + true_bias   # ⚠️ coef 是 float，x_line 是 numpy，无冲突

    plt.plot(x_line, y_true_line, color="green", label="真实直线 y = w*x + b")

    # 模型预测直线
    with torch.no_grad():
        y_pred_line = model(torch.tensor(x_line, dtype=torch.float32)).numpy()

    plt.plot(x_line, y_pred_line, color="red", label="模型预测直线")

    plt.legend()
    plt.title("线性回归拟合效果")
    plt.xlabel("x")
    plt.ylabel("y")
    plt.show()


if __name__ == '__main__':
    x ,y ,coef = create_dataset()
    print("="*40 + "查看数据" + "="*40)
    print(f"x:{x[:5]} \ny:{y[:5]}")  # 查看前5个数据

    print("="*40 + "开始训练" + "="*40)
    train(x,y,coef)
